import tensorflow as tf
import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
X = tf.constant([[1.0], [2.0]])
Y = tf.constant([[10.0], [20.0]])


class MyLinear(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(
            units=1,
            activation=None,
            kernel_initializer=tf.zeros_initializer(),
            bias_initializer=tf.zeros_initializer()
        )

    def call(self, input):
        out = self.dense(input)
        return out


model = MyLinear()
optim = tf.keras.optimizers.SGD(learning_rate=0.01)
for i in range(100):
    with tf.GradientTape() as tape:
        y_pred = model(X)
        loss = tf.reduce_mean(tf.square(y_pred - Y))
    grads = tape.gradient(loss, model.variables)

    optim.apply_gradients(grads_and_vars=zip(grads, model.variables))

# print(model.variables)
print(model.predict([3.0]))
